"""
这是使用ssh方式连接，下一个版本添加telnet连接。
2021/1/4
"""
import paramiko
import time
import os
import re
import csv
import config
import abc
import datetime
import telnetlib as tel
import threading
import logging
from openpyxl import Workbook
from openpyxl.chart import BarChart, LineChart, Reference

lock = threading.Lock()


class Switch(metaclass=abc.ABCMeta):
    """
    交换机原型。
    """
    def __init__(self, host, user, pwd):
        self.switch = self.connect(host, user, pwd)
        self.vender = None
        self.vender = self.detect_vender()
        # print('vender end')

    def detect_vender(self):
        """
        发现交换机厂家。
        """
        # system = None
        # 检测交换机厂家，ruijie或h3c。
        for sys in config.sys_detect_command.keys():
            sys_buffer = self.shell(config.sys_detect_command[sys])
            # print(config.sys_detect_regex[sys], sys_buffer)
            if re.search(config.sys_detect_regex[sys], sys_buffer):
                system = sys
                # print(system)
                if sys == 'ruijie':
                    self.commands = config.ruijie_command
                    self.patterns = config.ruijie_pattern
                    return sys
                elif sys == 'h3c':
                    self.commands = config.h3c_command
                    self.patterns = config.h3c_pattern
                    return sys

    @abc.abstractmethod
    def connect(self, host, user, pwd):
        """
        创建连接方式，telnet或ssh。
        """
        pass

    @abc.abstractmethod
    def shell(self, cmd):
        pass


class SSH2(Switch):
    def __init__(self, host, user, pwd):
        logging.info('%s %s', 'start SSH2()', host)

        super(SSH2, self).__init__(host, user, pwd)
        if self.vender == 'ruijie':
            self.switch.send('enable\n')
            self.switch.send('iflyun\n')

    def connect(self, host, user, pwd):
        logging.info('%s %s', 'start connect() in SSH2()', host)

        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(host, port=22, username=user, password=pwd, timeout=3)

        return client.invoke_shell()

    def shell(self, cmd):
        logging.info('%s %s', 'start shell() in SSH2()', cmd)

        self.switch.send(cmd + '\n')
        time.sleep(1)
        buffer = b''
        while True:
            once_buf = self.switch.recv(2000)
            buffer += once_buf
            # str_buffer = once_buf.decode('GBK')

            if b'--More--' in once_buf:  # 执行空格将继续接收数据。
                self.switch.send(' ')
                # time.sleep(1)
            elif once_buf.endswith(b'>') or once_buf.endswith(b'#'):  # 缓冲区数据结尾是>或#说明命令已经结束。
                break

        return buffer


class SSH:
    def __init__(self, host, user, pwd):
        self.client = paramiko.SSHClient()
        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.client.connect(host, port=22, username=user, password=pwd, timeout=3)
        self.switch = self.client.invoke_shell()
        self.commands = None
        self.patterns = None
        system = None

        # 检测交换机厂家，ruijie或h3c。
        for sys in config.sys_detect_command.keys():
            sys_buffer = self.shell(config.sys_detect_command[sys])
            if re.search(config.sys_detect_regex[sys], sys_buffer):
                system = sys
                if sys == 'ruijie':
                    self.commands = config.ruijie_command
                    self.patterns = config.ruijie_pattern
                elif sys == 'h3c':
                    self.commands = config.h3c_command
                    self.patterns = config.h3c_pattern
        # #^s+print(.+)

        if system == 'ruijie':
            self.switch.send('enable\n')
            self.switch.send('iflyun\n')

    def shell(self, cmd):
        self.switch.send(cmd + '\n')
        time.sleep(1)
        byte_buffer = b''
        while True:
            _buffer = self.switch.recv(2000)
            byte_buffer += _buffer
            str_buffer = _buffer.decode('GBK')

            if '--More--' in str_buffer:  # 执行空格将继续接收数据。
                self.switch.send(' ')
                # time.sleep(1)
            elif str_buffer.endswith('>') or str_buffer.endswith('#'):  # 缓冲区数据结尾是>或#说明命令已经结束。
                break

        return byte_buffer

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.client.close()


class Telnet2(Switch):
    def __init__(self, host, user, pwd):
        logging.info('%s %s', 'start in Telnet2()', host)
        self._encoding = 'utf-8'
        self.init_vender = True
        super(Telnet2, self).__init__(host, user, pwd)
        self.init_vender = False

    def connect(self, host, user, pwd) -> tel.Telnet:
        logging.info('%s %s', 'start connect() in Telnet2', host)
        login = [b'username', b'Username', b'UserName', b'login']
        password = [b'password', b'Password', b'PassWord']
        telnet = tel.Telnet(host)
        telnet.expect(login)
        telnet.write(bytes(user, encoding=self._encoding) + b'\n')
        telnet.expect(password)
        telnet.write(bytes(pwd, encoding=self._encoding) + b'\n')
        return telnet

    def shell(self, cmd):
        logging.info('%s %s', 'start shell in Telnet2', cmd)
        # print(cmd)
        self.switch.write(bytes(cmd, encoding=self._encoding) + b'\n')
        # self.switch.write(b'\n')

        count = 0
        buffer = b''
        while True:
            once_buf = self.switch.read_some()
            buffer += once_buf
            if b'--More--' in once_buf:
                # 对于多页信息需要发送空格来继续显示，对于大部分交换机来说可以发送b' \n'，但对于65来说这里只能使用b' '来发送，否则不能完整显示内容。
                self.switch.write(b' ')
            if self.vender and once_buf.endswith(b'#'):
                break
            if b'#' in once_buf:
                count += 1
            if count == 2 and once_buf.endswith(b'#'):
                break
        return buffer

    def test_shell(self, cmd):
        self.switch.write(bytes(cmd, encoding=self._encoding) + b'\n')
        buffer = b''
        self.switch.read_until(bytes(cmd))
        while True:
            once_buf = self.switch.read_some()
            buffer += once_buf
            if b'--More--' in once_buf:
                self.switch.write(b' \n')
            if once_buf.endswith(b'#'):
                break
        return buffer


class Telnet:
    def __init__(self, host, user, pwd):
        login = [b'username', b'Username', b'UserName']
        password = [b'password', b'Password', b'PassWord']
        self._encoding = 'utf-8'
        self.telnet = tel.Telnet(host)
        self.telnet.expect(login)
        self.telnet.write(bytes(user, encoding=self._encoding) + b'\n')
        self.telnet.expect(password)
        self.telnet.write(bytes(pwd, encoding=self._encoding) + b'\n')

    def shell(self, cmd):
        self.telnet.write(bytes(cmd, encoding=self._encoding) + b'\n')
        count = 0
        all_buffer = b''
        while True:
            raw = self.telnet.read_some()
            all_buffer += raw
            if b'--More--' in raw:
                self.telnet.write(b' \n')
            if b'#' in raw:
                count += 1
            if count == 2 and raw.endswith(b'#'):
                break

        return all_buffer


def _command(all_data, _config, all_cmd):
    """
    这是被多线程条用的方法，每一个多线程调用的都是每台设备执行完所有命令后得到的结果。
    """
    host = _config[0]
    conn = _config[1]
    user = _config[2]
    pwd = _config[3]
    logging.info('%s %s', 'start _command()', host)
    try:
        if conn == 'ssh':
            switch = SSH2(host, user, pwd)
        elif conn == 'telnet':
            switch = Telnet2(host, user, pwd)
    except Exception as e:
        # 如果ssh跟交换机连接失败，则进行空数据行填充。
        count = 9
        single_data = [host]
        for i in range(2, count):
            single_data += ['---']

        r = test_ping(host)
        if r:
            single_data += ['主机ping不通']
        else:
            single_data += ['其他原因导致无法连接']
        with lock:
            # all_data.append(single_data)
            all_data[host] = single_data
        exit()

    single_data = [host]
    for cmd in all_cmd:
        # 处理不同交换机使用不同命令去获取数据，通过在config.ambiguous_command变量中去指定一种数据存在多个命令。
        if cmd not in config.ambiguous_command:
            raw_data = switch.shell(switch.commands[cmd])
        else:
            # print(switch.commands)
            for _cmd in switch.commands[cmd]:
                raw_data = switch.shell(_cmd)
                if not re.search(config.invalid_key, raw_data):
                    break

        filter_data = re.findall(switch.patterns[cmd], raw_data.decode('GB2312'))

        # 这里对于cpu、mem、fan和temp都调用了单独的函数进行取值，clock、uptime取值放到公共的else中取值。
        if cmd == 'cpu' or cmd == 'mem':
            logging.info('%s %s', 'start cpu_mem() in _commnad()', host)
            single_data += cpu_mem(filter_data)
        elif cmd == 'fan':
            logging.info('%s %s', 'start fan() in _command()', host)
            r = fan(raw_data)
            single_data += r
        elif cmd == 'temp':
            logging.info('%s %s', 'start temp() in _command()', host)
            single_data += temp(raw_data)
        else:
            single_data += filter_data

    with lock:
        # all_data.append(single_data)
        all_data[host] = single_data


def commands(all_config: list, all_cmd: list):
    logging.info('%s', 'start commands')
    # all_data = []
    all_data = {}
    all_sorted_data = []
    titles = ['IP', '5sec', '1min', '5min', 'mem', 'fan', 'temp', 'clock', 'uptime']
    all_sorted_data.append(titles)

    # 这里对每一台设备使用了多线程。
    for _config in all_config:
        t = threading.Thread(target=_command, args=(all_data, _config, all_cmd, ))
        t.start()

    # 当所有交换机数据插入到all_data中时，再继续执行excel的生成。
    while True:
        print(len(all_data), len(all_config))
        if len(all_data) != len(all_config):
            # print(len(all_data), all_data.keys())
            time.sleep(1)
        else:
            break

    for host in [_host[0] for _host in all_config]:
        all_sorted_data.append(all_data[host])

    return all_sorted_data


def cpu_mem(filter_data):
    """
    过滤cpu和mem的方法。

    388: IndexError: list index out of range
    """
    logging.info('%s %s', 'start in cpu_mem()', filter_data)

    convert_data = []
    for _data in filter_data:
        convert_data.append(float(_data.split('%')[0]))

    try:
        if len(filter_data) > 3:
            grouping_data = [[convert_data[i], convert_data[i + 1], convert_data[i + 2]] for i in
                             range(1, len(convert_data)) if i % 3 == 0]
            _max = 0
            for min_index in range(1, len(grouping_data)):
                if sum(grouping_data[_max]) < sum(grouping_data[min_index]):
                    _max = min_index
            return grouping_data[_max]
        else:
            return convert_data
    except Exception as e:
        logging.info('%s %s', 'exception in cpu_mem()', e, exc_info=True)


def get_row(lines, status_patterns):
    """
    获取多行命令结果的行数。
    """
    row = None
    # 定位温度文字所在输出行。
    for n, line in enumerate(lines):
        for match in status_patterns:
            if re.search(match, line):
                row = n
                break
        if row:
            break

    return row


def get_column(lines, row, status_patterns):
    """
    获取多行数据的列数。
    """
    columns = re.findall(b'\w+  +|\w+$', lines[row])
    for i, _column in enumerate(columns):
        for pattern in status_patterns:
            if pattern in _column:
                return i


def get_column_region(lines, row):
    """
    将多行数据中状态栏那行数据分割成多列区间数据。
    """
    column_region = [0]
    split_words = re.findall(b'(\w+[ -_]\w+|\w+) *', lines[row])
    # split_words = re.findall(b'(\w+|\w+ \w+) *', lines[row])
    for word in split_words[1:]:
        index = re.search(re.escape(word), lines[row]).start() - 1
        column_region.append(index)

    return column_region


def find_data(lines, status_patterns: list):
    """
    针对某些命令返回多行结果进行查找状态所在行和列。
    return: 一个关键字所在行，一个索引位置。
    """
    row = None  # title_match所在行。
    column = None  # title_match所在列。
    indices = []
    # 定位温度文字所在输出行。
    for n, line in enumerate(lines):
        for match in status_patterns:
            if re.search(match, line):
                row = n
                break
        if row:
            break

    # 记录包含温度文字那一行的每列文字起始位置。
    space_count = 0
    indices.append(0)
    for n, char in enumerate(lines[row]):
        if chr(char) == ' ':
            space_count += 1
            continue
        if re.search('[a-zA-Z]', chr(char)) and space_count > 1:
            indices.append(n)
            space_count = 0

    for pattern in status_patterns:
        if not re.search(pattern, lines[row]):
            continue
        for i in range(len(indices)):
            txt = lines[row][indices[i]:] if len(indices) == i + 1 else lines[row][
                                                                                indices[i]:indices[i + 1]]
            if re.search(pattern, txt):
                column = i
                break
        if column:
            break

    return indices


def fan(raw_data):
    """
    对fan2按照temp的实现方式去实现。
    """
    logging.info('%s %s', 'start in fan()', raw_data)

    fan_status = b'( ok| *normal| *Normal| on| N/A)'
    fan_state_priority = {
        b'ok': 1,
        b'OK': 1,
        b'normal': 2,
        b'Normal': 2,
        b'on': 3,
        b'ON': 3,
        b'n/a': 9,
        b'N/A': 9
    }
    single_line = b'(State|state).+' + fan_status
    multiple_line = [b'status', b'Status', b'State', b'state']
    lines = [_.replace(b'\r', b'') for _ in raw_data.split(b'\n')]
    _base_line = re.findall(fan_status, raw_data)
    if len(_base_line) == 1:
        base_line = [_base_line[0].replace(b' ', b'')]
        return base_line
    elif len(_base_line) > 1:
        max_status = None
        single_results = re.findall(single_line, raw_data)
        # 在同一行同时拥有”state“和风扇结果，直接过滤出结果，否则执行多行风扇数据的过滤处理。
        if single_results:
            for line in single_results:
                temp_status = (line[1], fan_state_priority[line[1]])
                if not max_status:
                    max_status = temp_status
                    continue
                elif max_status[1] > temp_status[1]:
                    max_status = temp_status
            if max_status:
                return [max_status[0]]
        else:
            row = get_row(lines, multiple_line)
            column = get_column(lines, row, multiple_line)
            column_region = get_column_region(lines, row)
            for line in lines[row:]:
                if not re.search(fan_status, line):
                    continue
                raw_fan_result = line[column_region[column]:] if len(column_region) == column + 1 else line[column_region[column]: column_region[column + 1]]
                fan_result = raw_fan_result.replace(b' ', b'')
                if not max_status:
                    max_status = (fan_result, fan_state_priority[fan_result])
                temp_status = (fan_result, fan_state_priority[fan_result])
                if max_status[1] < temp_status[1]:
                    max_status = temp_status
            return [max_status[0]]


def temp(raw_data):
    """
    获取温度过滤数据。
    """
    logging.info('%s %s', 'start in temp()', raw_data)

    max_temp = -1
    status_patterns = [b'current|Current', b'Temperature']
    column_region = []
    # word_match = [b'Temp', b'Temperature']
    str_match = b'Temp|Temperature|Measure'
    num_match = b'[1-9][0-9].[0-9]{3}|[1-9][0-9]'
    temp_count = []
    lines = [_.replace(b'\r', b'') for _ in raw_data.split(b'\n')]

    # 一行内同时拥有温度单词和温度数字，则保存当前行内温度。
    for line in lines:
        if re.search(str_match, line) and re.search(num_match, line):
            for i in re.findall(num_match, line):
                # temp_count.append(float(i))
                if float(i) > max_temp:
                    max_temp = float(i)

    if max_temp != -1:
        return [max_temp]

    # 当遇到核心交换机有多个板卡时，从这些办卡中获取温度最高的数值。
    if len(temp_count) == 0:
        row = None  # title_match所在行。
        column = None   # title_match所在列。
        # 定位温度文字所在输出行。
        for n, line in enumerate(lines):
            for match in status_patterns:
                if re.search(match, line):
                    row = n
                    break
            if row:
                break

        # 记录包含温度文字那一行的每列文字起始位置。
        space_count = 0
        column_region.append(0)
        for n, char in enumerate(lines[row]):
            if chr(char) == ' ':
                space_count += 1
                continue
            if re.search('[a-zA-Z]', chr(char)) and space_count > 1:
                column_region.append(n)
                space_count = 0

        for pattern in status_patterns:
            if not re.search(pattern, lines[row]):
                continue
            for i in range(len(column_region)):
                txt = lines[row][column_region[i]:] if len(column_region) == i + 1 else lines[row][column_region[i]:column_region[i+1]]
                if re.search(pattern, txt):
                    column = i
                    break
            if column:
                break

        # 从包含温度文字列之后每一列中查看温度数字。
        for line in lines[row:]:
            if not re.search(b' [0-9]{2} ', line):
                continue
            temp_txt = line[column_region[column]:] if len(column_region) == column + 1 else line[column_region[column]: column_region[column + 1]]
            # print(temp_txt)
            match_temp = re.findall(b'[0-9]{2}', temp_txt)
            if not match_temp:
                continue
            for temp_value in match_temp:
                if int(temp_value) > max_temp:
                    max_temp = int(temp_value)

        return [max_temp]


def is_contain(raw_data, keys):
    """
    验证关键字，对执行命令返回的信息验证是否包含只当关键字。
    """
    if isinstance(keys, bytes):
        return True if keys in raw_data else False
    elif isinstance(keys, list):
        for k in keys:
            if k in raw_data:
                return True
        else:
            return False


def test_ping(host):
    """
    当主机连接超时时测试主机是否ping通。
    """
    cmd = 'ping -w 1000 -n 1 {} 1> nul'.format(host)
    results = []
    for _ in range(3):
        results.append(os.system(cmd))
    return all(results)


def get_config(path):
    """
    从csv文件中获取配置信息。
    """
    config_file = []
    with open(path) as file:
        csv_file = csv.reader(file)
        for line in csv_file:
            config_file.append(line)

    return config_file


def create_bar_chart(sheet_data, d_range):
    """
    创建excel中展示的柱状图表。
    """
    chart = BarChart()
    data = Reference(sheet_data, min_row=1, max_row=sheet_data.max_row, min_col=d_range[0], max_col=d_range[1])
    label = Reference(sheet_data, min_col=1, min_row=2, max_row=sheet_data.max_row)
    chart.add_data(data, titles_from_data=True)
    chart.set_categories(label)

    return chart


def create_line_chart(sheet_data, d_range):
    """
    创建excel中展示的线性图标。
    """
    chart = LineChart()
    data = Reference(sheet_data, min_row=1, max_row=sheet_data.max_row, min_col=d_range[0], max_col=d_range[1])
    label = Reference(sheet_data, min_col=1, min_row=2, max_row=sheet_data.max_row)
    chart.add_data(data, titles_from_data=True)
    chart.set_categories(label)

    return chart


def create_sheet(file_name, data):
    """
    创建excel中sheet工作表。
    """
    data_range = [(2, 4), (5, 5), (7, 7)]
    choice_chart = {
        0: create_bar_chart,
        1: create_line_chart,
        2: create_line_chart,
    }
    chart_position = ['k1', 'k16', 'k32']
    labels = ['CPU使用率', '内存使用率', '温度']
    workbook = Workbook()
    sheet = workbook.create_sheet('test', 0)
    for d in data:
        sheet.append(d)
    for n, _range in enumerate(data_range):
        chart = choice_chart[n](sheet, _range)
        chart.title = labels[n]
        sheet.add_chart(chart, chart_position[n])

    workbook.save(file_name)


def get_csv(path):
    data = []
    with open(path) as file:
        f = csv.reader(file)
        for i, line in enumerate(f):
            if i > 0:
                for j, _line in enumerate(line):
                    if j in [1, 2, 3, 4, 6]:
                        line[j] = float(_line)

            data.append(line)
    return data


def main():
    logging.basicConfig(filename='g:\\switches-running.log', filemode='w', level='INFO', format='%(asctime)s %(message)s', datefmt='%H:%M:%S')
    start = time.time()
    date = datetime.datetime.now()
    path = 'g:\\'
    filename = '云计算网络设备日常巡检表' + str(date.year) + str(date.month) + str(date.day) + '.xlsx'
    all_config = get_config(path + 'config.csv')
    data = commands(all_config, ['cpu', 'mem', 'fan', 'temp', 'clock', 'uptime'])

    # 使用csv数据测试图形。
    # data = get_csv('g:\\data.csv')

    create_sheet(path + filename, data)
    print('time-consuming {}s'.format(str(round(time.time() - start, 2))))


if __name__ == '__main__':
    main()